import torch
import torch.nn as nn

class SymmetricQuantizer(nn.Module):
    def __init__(self):
        super(SymmetricQuantizer, self).__init__()
        self.qmin = -128  # int8 的最小值
        self.qmax = 127   # int8 的最大值

    def forward(self, x):
        # 1. 截断到 [-1, 1]
        x = torch.clamp(x, -1.0, 1.0)

        # 2. 计算量化参数 scale
        scale = (self.qmax - self.qmin) / 2.0  # 对称映射比例
        zero_point = 0  # 对称量化的 zero_point 为 0

        # 3. 量化
        x_quantized = torch.round(x * scale).to(torch.int8)

        # 4. 反量化
        x_dequantized = x_quantized.float() / scale

        return x_quantized, x_dequantized


# 测试
input_tensor = torch.rand(10) * 2 - 1
quantizer = SymmetricQuantizer()
x_quantized, x_dequantized = quantizer(input_tensor)

print("原始输入:", input_tensor)
print("量化后的值:", x_quantized)
print("反量化后的值:", x_dequantized)


def calculate_mse(original_tensor, quantized_tensor):
    mse = torch.mean((original_tensor - quantized_tensor) ** 2).item()
    return mse

mse = calculate_mse(input_tensor, x_dequantized)
print("MSE:", mse)